#include "defines.h"
#include "lib.h"
#include "kozos.h"
#include "nicdrv.h"
#include "netdef.h"
#include "netlib.h"
#include "arp.h"

#define ARP_CACHE_NUM   8
#define ARP_TIMEOUT     600
#define ARP_RETRY       5
#define QUEUE_MAX       4

typedef struct _packet_q {
    union {
        struct _packet_q *next;
        uint8 IP[4];
    } un;
    uint8 *packet;
    int size;
} packet_q;

typedef struct {
#define ARP_FREE        0
#define ARP_RESOLVED    1
#define ARP_WAIT        2
    int state;
    uint8 IP[4];
    uint8 MAC[6];
    int count;
    int qsize;
    packet_q *head;
    packet_q *tail;
} arp_cache_entry;

static arp_cache_entry Acache[ARP_CACHE_NUM];
static kz_msgbox_id_t arp_mbox = MSGBOX_NULL;
static uint8 src_IP[] = SRC_IP;

void arp_init(void)
{
    ASSERT(memeq(Acache, 0, sizeof(Acache)), "BSS isn't cleared");
    arp_mbox = kz_getmsgbox();
}

void arp_setIP(const uint8 IP[])
{
    uint8 *new_IP;
    ASSERT(arp_mbox != MSGBOX_NULL, "no arp_mbox");
    new_IP = kz_kmalloc(4);
    memcpy(new_IP, IP, 4);
    kz_send(arp_mbox, 'I', (char *)new_IP, 0, 0);
}

void arp_handle(uint8 *buf)
{
    ASSERT(arp_mbox != MSGBOX_NULL, "no arp_mbox");
    if (kz_send(arp_mbox, 'h', (char *)buf, 1, NET_MBOXSIZE) == -1)
        kz_kmfree(buf);
}

void arp_resolve(const uint8 IP[], uint8 *buf, int size)
{
    packet_q *p;

    ASSERT(arp_mbox != MSGBOX_NULL, "no arp_mbox");
    p = kz_kmalloc(sizeof(packet_q));
    memcpy(p->un.IP, IP, 4);
    p->packet = buf;
    p->size = size;
    if (kz_send(arp_mbox, 'r', (char *)p, 1, NET_MBOXSIZE) == -1) {
        kz_kmfree(buf);
        kz_kmfree(p);
    }
}

void arp_timer()
{
    ASSERT(arp_mbox != MSGBOX_NULL, "no arp_mbox");
    kz_send(arp_mbox, 't', NULL, 2, 1);
}

static int check_arp_packet(arp_packet *ap)
{
    return ap->arp_hardware_type == 1 &&
           ap->arp_protocol_type == 0x0800 &&
           ap->arp_hardware_length == 6 &&
           ap->arp_protocol_length == 4 &&
           (ap->arp_operation == 1 || ap->arp_operation == 2) &&
           !memcmp(ap->arp_dst_IP, src_IP, 4);
}

static void arp_request(uint8 *IP)
{
    ethernet_header *ehdr = (ethernet_header *)kz_kmalloc(42);
    arp_packet *apacket = (arp_packet *)(ehdr + 1);

    apacket->arp_hardware_type = 1;
    apacket->arp_protocol_type = 0x0800;
    apacket->arp_hardware_length = 6;
    apacket->arp_protocol_length = 4;
    apacket->arp_operation = 1;
    nic_MAC(apacket->arp_src_MAC);
    memset(apacket->arp_dst_MAC, 0, 6);
    memcpy(apacket->arp_src_IP, src_IP, 4);
    memcpy(apacket->arp_dst_IP, IP, 4);

    memset(ehdr->eth_dst_MAC, 0xff, 6);
    memcpy(ehdr->eth_src_MAC, apacket->arp_src_MAC, 6);
    ehdr->eth_ethernet_type = 0x0806;

    nic_send((uint8 *)ehdr, 42);
}

static void arp_reply(uint8 *buffer) // bufferはチェック済ARP要求フレーム
{
    ethernet_header *ehdr = (ethernet_header *)buffer;
    arp_packet *apacket = (arp_packet *)(ehdr + 1);

    apacket->arp_operation = 2;
    memcpy(apacket->arp_dst_MAC, apacket->arp_src_MAC, 6);
    nic_MAC(apacket->arp_src_MAC);
    memcpy(apacket->arp_dst_IP, apacket->arp_src_IP, 4);
    memcpy(apacket->arp_src_IP, src_IP, 4);

    memcpy(ehdr->eth_dst_MAC, apacket->arp_dst_MAC, 6);
    memcpy(ehdr->eth_src_MAC, apacket->arp_src_MAC, 6);
    ehdr->eth_ethernet_type = 0x0806;

    nic_send(buffer, 42);
}

static void cache_free(arp_cache_entry *p)
{
    packet_q *q, *r;

    ASSERT(p->state != ARP_FREE, "cache_free: already free");

    for (q = p->head; q; ) {
        r = q->un.next;
        kz_kmfree(q->packet);
        kz_kmfree(q);
        q = r;
    }

    memset(p, 0, sizeof(*p));
}

static arp_cache_entry *cache_alloc()
{
    static arp_cache_entry *remove_last = Acache;

    if (++remove_last == Acache + ARP_CACHE_NUM)
        remove_last = Acache;
    cache_free(remove_last);

    return remove_last;
}

static void cache_add(uint8 IP[], uint8 MAC[])
{
    arp_cache_entry *p, *freep = NULL;

    for (p = Acache; p < Acache + ARP_CACHE_NUM; p++) {
        if (p->state == ARP_FREE) {
            if (!freep) freep = p;
        } else {
            if (!memcmp(p->IP, IP, 4)) break;
        }
    }

    if (p < Acache + ARP_CACHE_NUM) {

        if (p->state == ARP_RESOLVED) {

            memcpy(p->MAC, MAC, 6);
            p->count = ARP_TIMEOUT;

        } else {

            ASSERT(p->state == ARP_WAIT, "Bad state");

            p->state = ARP_RESOLVED;
            memcpy(p->MAC, MAC, 6);
            p->count = ARP_TIMEOUT;

            {
                packet_q *q, *r;

                for (q = p->head; q; ) {
                    r = q->un.next;
                    memcpy(q->packet, MAC, 6);
                    nic_send(q->packet, q->size);
                    kz_kmfree(q);
                    q = r;
                }

                p->head = p->tail = NULL;
                p->qsize = 0;
            }
        }

    } else {

        if (!freep) {
            freep = cache_alloc();
        }

        freep->state = ARP_RESOLVED;
        memcpy(freep->IP, IP, 4);
        memcpy(freep->MAC, MAC, 6);
        freep->count = ARP_TIMEOUT;

    }
}

static void cache_resolve(packet_q *param)
{
    arp_cache_entry *p, *freep = NULL;

    for (p = Acache; p < Acache + ARP_CACHE_NUM; p++) {
        if (p->state == ARP_FREE) {
            if (!freep) freep = p;
        } else {
            if (!memcmp(p->IP, param->un.IP, 4)) break;
        }
    }

    if (p < Acache + ARP_CACHE_NUM) {

        if (p->state == ARP_RESOLVED) {

            memcpy(param->packet, p->MAC, 6);
            nic_send(param->packet, param->size);
            kz_kmfree(param);

        } else {

            ASSERT(p->state == ARP_WAIT, "Bad state");

            if (p->qsize < QUEUE_MAX) {

                param->un.next = NULL;

                if (p->head) {
                    p->tail->un.next = param;
                } else {
                    p->head = param;
                }
                p->tail = param;
                p->qsize++;

            } else {
                kz_kmfree(param->packet);
                kz_kmfree(param);
                puts("arp queue is full");
            }
        }

    } else {

        arp_request(param->un.IP);

        if (!freep) {
            freep = cache_alloc();
        }

        freep->state = ARP_WAIT;
        memcpy(freep->IP, param->un.IP, 4);
        freep->count = ARP_RETRY;

        param->un.next = NULL;
        freep->head = freep->tail = param;
        freep->qsize = 1;

    }
}

void arp_print(void)
{
    int i;

    puts("# IP MAC state count qsize\n");

    for (i = 0; i < ARP_CACHE_NUM; i++) {
        arp_cache_entry *p = &Acache[i];
        if (p->state != ARP_FREE) {
            putval(i, 0);
            puts(" ");
            print_IP(p->IP);
            puts(" ");
            print_MAC(p->MAC);
            puts(" ");
            putval(p->state, 0);
            puts(" ");
            putval(p->count, 0);
            puts(" ");
            putval(p->qsize, 0);
            puts("\n");
        }
    }
}

void arp_main(int argc, char *argv)
{
    int cmd;
    char *p;

    while (1) {
        kz_recv(arp_mbox, &cmd, &p);

        if (cmd == 'r') {

            cache_resolve((packet_q *)p);

        } else if (cmd == 't') {
            arp_cache_entry *e;

            for (e = Acache; e < Acache + ARP_CACHE_NUM; e++) {
                if (e->state == ARP_RESOLVED) {

                    if (e->count == 0) {
                        cache_free(e);
                    } else {
                        e->count--;
                    }

                } else if (e->state == ARP_WAIT) {

                    if (e->count == 0) {
                        cache_free(e);
                    } else {
                        arp_request(e->IP);
                        e->count--;
                    }
                }
            }

        } else if (cmd == 'h') {

            arp_packet *ap = (arp_packet *)(p + sizeof(ethernet_header));

            if (check_arp_packet(ap)) {
                cache_add(ap->arp_src_IP, ap->arp_src_MAC);
                if (ap->arp_operation == 1) {
                    arp_reply((uint8 *)p);
                    continue;
                }
            }

            kz_kmfree(p);

        } else if (cmd == 'I') {

            memcpy(src_IP, p, 4);
            kz_kmfree(p);

        } else {
            ASSERT(0, "arp: Bad cmd");
        }
    }
}
